package clickhouse

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"

	"gitlab.bangdao-tech.com/ops/ck-migrator/pkg/logger"

	"gitlab.bangdao-tech.com/ops/ck-migrator/pkg/config"

	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"

	"github.com/riete/errors"

	"github.com/ClickHouse/clickhouse-go/v2"
)

type ExecParams struct {
	Sql              string
	Failed           chan string
	Total            int
	DBTablePartition string
	Result           chan string
	OriSqlString     string
	ShardInfo        string // 添加分片信息字段
	IsLastShard      bool   // 是否是最后一个分片
}

type Clickhouse struct {
	config   config.ClickHouse
	conn     clickhouse.Conn
	progress map[string]*atomic.Uint64
	l        sync.Mutex
}

func (c *Clickhouse) Open() errors.Error {
	option := &clickhouse.Options{
		Auth: clickhouse.Auth{
			Username: c.config.Username,
			Password: c.config.Password,
		},
		Addr: []string{fmt.Sprintf("%s:%d", c.config.Host, c.config.Port)},
	}
	var err error
	c.conn, err = clickhouse.Open(option)
	return errors.NewFromErr(err)
}

func (c *Clickhouse) Exec(p *ExecParams) {
	c.l.Lock()
	writeRows, ok := c.progress[p.DBTablePartition]
	if !ok {
		writeRows = new(atomic.Uint64)
		c.progress[p.DBTablePartition] = writeRows
	}
	c.l.Unlock()
	
	// 跟踪是否已经发送了完成状态
	completionSent := false
	
	// 跟踪最后一次发送的进度信息，避免重复发送相同的进度
	lastProgressInfo := ""
	
	// 定义发送进度信息的函数
	sendProgress := func(wroteRows uint64) {
		wrote := int(wroteRows / 2)
		var progressInfo string
		
		if wrote >= p.Total { // 使用 >= 而不是 == 来确保完成状态被正确标记
			// 避免重复发送完成状态
			if !completionSent {
				if p.ShardInfo != "" {
					// 只有最后一个分片才显示"迁移完成"
					if p.IsLastShard {
						progressInfo = fmt.Sprintf("Partition: %s | Slice: %s | Progress: %d/%d, 迁移完成", p.DBTablePartition, p.ShardInfo, wroteRows/2, p.Total)
						completionSent = true
					} else {
						progressInfo = fmt.Sprintf("Partition: %s | Slice: %s | Progress: %d/%d", p.DBTablePartition, p.ShardInfo, wroteRows/2, p.Total)
					}
				} else {
					progressInfo = fmt.Sprintf("Partition: %s | Progress: %d/%d, 迁移完成", p.DBTablePartition, wroteRows/2, p.Total)
					completionSent = true
				}
			}
		} else {
			if p.ShardInfo != "" {
				progressInfo = fmt.Sprintf("Partition: %s | Slice: %s | Progress: %d/%d", p.DBTablePartition, p.ShardInfo, wroteRows/2, p.Total)
			} else {
				progressInfo = fmt.Sprintf("Partition: %s | Progress: %d/%d", p.DBTablePartition, wroteRows/2, p.Total)
			}
		}
		
		// 避免重复发送相同的进度信息
		if progressInfo != "" && progressInfo != lastProgressInfo {
			p.Result <- progressInfo
			lastProgressInfo = progressInfo
		}
	}
	
	// 定义发送完成状态的函数
	sendCompletion := func() {
		// 如果完成状态还没有发送过，则发送
		if !completionSent {
			var progressInfo string
			if p.ShardInfo != "" {
				// 只有最后一个分片才显示"迁移完成"
				if p.IsLastShard {
					progressInfo = fmt.Sprintf("Partition: %s | Slice: %s | Progress: %d/%d, 迁移完成", p.DBTablePartition, p.ShardInfo, p.Total, p.Total)
				} else {
					// 非最后一个分片不发送完成状态
					progressInfo = fmt.Sprintf("Partition: %s | Slice: %s | Progress: %d/%d", p.DBTablePartition, p.ShardInfo, p.Total, p.Total)
				}
			} else {
				// 对于没有分片的任务，发送完成状态
				progressInfo = fmt.Sprintf("Partition: %s | Progress: %d/%d, 迁移完成", p.DBTablePartition, p.Total, p.Total)
			}
			
			// 避免重复发送相同的进度信息
			if progressInfo != "" && progressInfo != lastProgressInfo {
				p.Result <- progressInfo
				lastProgressInfo = progressInfo
			}
		}
	}
	
	ctx := clickhouse.Context(context.Background(), clickhouse.WithProgress(func(progress *clickhouse.Progress) {
		writeRows.Add(progress.WroteRows)
		sendProgress(writeRows.Load())
	}))
	
	if err := c.conn.Exec(ctx, p.Sql); err != nil {
		logger.Errorf("执行 %s 失败, %s", p.Sql, err.Error())
		p.Failed <- p.OriSqlString
	} else {
		// SQL执行完成，发送完成状态
		sendCompletion()
	}
}

func (c *Clickhouse) Query(ctx context.Context, query string, args ...any) (driver.Rows, errors.Error) {
	rows, err := c.conn.Query(ctx, query, args...)
	return rows, errors.NewFromErr(err)
}

func (c *Clickhouse) QueryRow(ctx context.Context, query string, args ...any) driver.Row {
	return c.conn.QueryRow(ctx, query, args...)
}

func (c *Clickhouse) Select(ctx context.Context, dest any, query string, args ...any) errors.Error {
	return errors.NewFromErr(c.conn.Select(ctx, dest, query, args...))
}

func New(config config.ClickHouse) (*Clickhouse, errors.Error) {
	c := &Clickhouse{config: config, progress: make(map[string]*atomic.Uint64)}
	if err := c.Open(); err != nil {
		return nil, err.Trace("连接数据库失败")
	}
	return c, nil
}